import pathlib
import sys
import json
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models.yi import ChatYi
import rclpy
from rclpy.executors import ExternalShutdownException
from geometry_msgs.msg import PoseStamped
from semantic_nav.navigator import Navigator
from semantic_nav.llm_utils import ChatChain
from semantic_nav.slam_utils import convert_dict_to_Pose
from semantic_nav.log_utils import get_logger
logger = get_logger()


class SemanticMap:
    def __init__(self, waypoint_info_paths):
        with open(waypoint_info_paths, "r") as f:
            self.waypoints_info = json.load(f)

    def get_object_pose(self, object_name):
        """
        get the pose of the object
        """
        for scene, objects in self.waypoints_info.items():
            if object_name in objects:
                return objects[object_name]

        return None

    def load_waypoint_info(self):
        """
        load the waypoint information
        """
        return "{}".format(self.waypoints_info)


class LlmNavigatorInMap(Navigator):
    def __init__(self, map_path=None) -> None:
        super().__init__()
        self.semantic_map = SemanticMap(map_path)  # put the path of the waypoint information file here
        logger.info(self.semantic_map.load_waypoint_info())
        system_template = "You are a navigation system for robotics. You need to analysis which object waypoint should the robot go to according to the user's input and the waypoints information. The waypoints information are organized in json format, the first level is scene, and the second level is object waypoint. Only output the name of the object that best matches the user's input to me, no other things."
        user_template = "The information of the waypoints are: {waypoint_info}"
        prompt_template = ChatPromptTemplate.from_messages([
            ('system', system_template),
            ('user', user_template),
            ('user', '{task}')
        ])
        partial_template = prompt_template.partial(waypoint_info=self.semantic_map.load_waypoint_info())
        self.chain = ChatChain(prompt_template=partial_template, model=ChatYi(model="yi-large", temperature=0.5))

    def navigate_to_object(self, response_str):
        """
        find the corresponding pose
        use move base to navigate
        """
        pose_dict = self.semantic_map.get_object_pose(response_str)
        if pose_dict is not None:
            # set the goal pose
            pose_stamped = PoseStamped()
            pose_stamped.header.frame_id = "map"
            pose_stamped.pose = convert_dict_to_Pose(pose_dict)
            logger.info("Going to: {}".format(response_str))
            logger.info(pose_stamped)
            # send goal to navigation2 action server
            self.navigate([pose_stamped])
            return True

        logger.info(f"{response_str}\nObject not found")
        return False

    def invoke(self, command):
        response = self.chain.invoke({"task": command})
        return self.navigate_to_object(response)


def main(args=None):
    rclpy.init(args=args)
    try:
        from ament_index_python.packages import get_package_prefix
        package_path = get_package_prefix('semantic_nav')
        map_path = pathlib.Path(package_path).joinpath('json/map.json')
        if not map_path.exists():
            logger.error("Map file not found")
            return
        llm_navigator = LlmNavigatorInMap(str(map_path))
        while rclpy.ok():
            user_input = input("User: ")
            if user_input.lower() in ["exit", "quit", "bye"]:
                logger.info("Chat session exit!")
                break
            llm_navigator.invoke(user_input)
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        sys.exit(1)


if __name__ == "__main__":
    main()
